﻿// Copyright (c) 2014 Morten Bakkedal
// This code is published under the MIT License.

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Text;
using System.Reflection;
using System.CodeDom.Compiler;
using Microsoft.CSharp;

using FuncLib.DualFunctions.Collections;
using FuncLib.DualFunctions.Expressions;
using FuncLib.Functions;

namespace FuncLib.DualFunctions.Compilation
{
	internal class DualCodeGenerator
	{
		private Dictionary<string, string> methods = new Dictionary<string, string>();
		private List<string> statements = new List<string>();
		private List<IndexList> offsetsSorted = new List<IndexList>();
		private Dictionary<IndexList, int> offsetsIndices = new Dictionary<IndexList, int>();
		//private HashedList<IndexList> indexLists = new HashedList<IndexList>();

		private List<GeneratedFunction> compiledFunctions;
		private Dictionary<DualFunction, GeneratedFunction> compiledMapping;
		private string generatedCode, identifier;
		private uint nextMethod;
		private int nextStorage;
		private bool useRnd = false;
		private Random rnd = new Random(2);
		private Stopwatch sw_unary = new Stopwatch();

		private StorageAllocator storageAllocator = new StorageAllocator();

		/*public DualCodeGenerator()
		{
			compiledMapping = new Dictionary<DualFunction, GeneratedFunction>();
			compiledFunctions = new List<GeneratedFunction>();
		}

		public void Compile(List<DualFunction> sortedFunctions, Dictionary<DualFunction, List<DualFunction>> outOfScopeMap)
		{
			foreach (DualFunction function in sortedFunctions)
			{
				function.Compile(this); // don't use return values

				if (outOfScopeMap.ContainsKey(function))
				{
					foreach (DualFunction unusedFunction in outOfScopeMap[function])
					{
						storageAllocator.Free(compiledMapping[unusedFunction].FunctionValueStorage);
						storageAllocator.Free(compiledMapping[unusedFunction].GradientStorage);
						storageAllocator.Free(compiledMapping[unusedFunction].HessianStorage);

						// Make sure that the derivative structure can be cleaned by the GC.
						compiledMapping.Remove(unusedFunction);
					}
				}
			}
		}

		//public DualCodeGenerator(StorageAllocator storageAllocator)
		//{
		//    this.storageAllocator = storageAllocator;
		//}*/

		public DualCodeGenerator(DualFunction function)
		{
			compiledFunctions = new List<GeneratedFunction>();
			compiledMapping = new Dictionary<DualFunction, GeneratedFunction>();

			generatedCode = GenerateCode(function, out identifier);
		}

		public int AddOffsets(IndexList offsets)
		{
			int index;
			if (!offsetsIndices.TryGetValue(offsets, out index))
			{
				index = offsetsSorted.Count;
				offsetsSorted.Add(offsets);
				offsetsIndices.Add(offsets, index);
			}

			return index;
		}

		/*public int AddIndexList(IndexList indexList)
		{
			int index = indexLists.IndexOf(indexList);
			if (index == -1)
			{
				index = indexLists.Count;
				indexLists.Add(indexList);
			}

			return index;
		}*/

		public GeneratedFunction GetCompiledFunction(DualFunction function)
		{
			GeneratedFunction compiledFunction;
			if (!compiledMapping.TryGetValue(function, out compiledFunction))
			{
				// Generate code for this particular function. Store the method body as a string with some additional properties.
				compiledFunction = function.Compile(this);

				compiledFunctions.Add(compiledFunction);
				compiledMapping.Add(function, compiledFunction);
			}

			return compiledFunction;
		}

		public Storage GetStorage(int length)
		{
			int startIndex = nextStorage;
			nextStorage += length + (useRnd ? rnd.Next(4) : 0);
			return new Storage(startIndex, length);
		}

		public FunctionValueStorage GetFunctionValueStorage()
		{
			int startIndex = nextStorage;
			nextStorage += 1 + (useRnd ? rnd.Next(3) : 0);
			return new FunctionValueStorage(startIndex);
			//return new FunctionValueStorage(nextStorage++);
		}

		private string GenerateCode(DualFunction function, out string identifier)
		{
			return GenerateCodeNEW(function, out identifier);
		}

		/// <summary>
		/// Creates a new compiled function defined by a variable (i.e. the identity function).
		/// </summary>
		public GeneratedFunction AddVariable(DualVariable variable)
		{
			FunctionBody body = new FunctionBody(this, "g");

			// Build gradient and Hessian.
			SortedGradient gradient = new SortedGradient();
			SortedHessian hessian = new SortedHessian();

			gradient.AddOnce(variable);

			FunctionValueStorage functionValueStorage = GetFunctionValueStorage();
			Storage gradientStorage = GetStorage(1);
			Storage hessianStorage = GetStorage(0);

			body.AddFunctionValue(new VariableExpression(variable));
			body.AddValues(functionValueStorage, gradientStorage, new Expression[] { 1.0 });

			return AddCompiledFunction(variable, body, functionValueStorage, gradient, gradientStorage, hessian, hessianStorage);
		}

		public GeneratedFunction AddUnaryFunction(DualFunction function, GeneratedFunction f, Expression g, Expression g1, Expression g11)
		{
			////body.ResultFunction
			//FunctionExpression expression = new FunctionExpression();

			//// Link the inner function to its identifier in this method.
			//Dictionary<FunctionExpression, string> identifiers = new Dictionary<FunctionExpression, string>();
			//identifiers[expression] = "g";
			//identifiers[f.ExpressionNEW] = "f";

			f.Expression.Link("f");

			FunctionBody body = new FunctionBody(this, "g", new FunctionParameter("f", f));

			// Replace by references to local variables.
			g1 = Expression.MarkLocal(g1);
			g11 = Expression.MarkLocal(g11);

			// List of expressions needed to compute the function value and derivatives.
			HashedList<Expression> expressions = new HashedList<Expression>();
			expressions.Add(g);

			// Build gradient.
			SortedGradient gradient = new SortedGradient();

			// And output and intermediate Hessians.
			SortedHessian hessian = new SortedHessian();
			SortedHessian hessian1 = new SortedHessian();
			SortedHessian hessian11 = new SortedHessian();

			if (f.Hessian.Count == f.Gradient.HessianCount)
			{
				if (!g1.IsZero())
				{
					expressions.AddOnce(g1);

					gradient = f.Gradient;
					hessian = f.Hessian;
					hessian1 = f.Hessian;
				}

				if (!g11.IsZero())
				{
					expressions.AddOnce(g11);

					hessian = f.Hessian;
					hessian11 = f.Hessian;
				}
			}
			else
			{
				IList<DualVariable> variables = f.Gradient;

				for (int i = 0; i < variables.Count; i++)
				{
					DualVariable variable1 = variables[i];

					if (!DualVariable.CheckFirstOrder(variable1))
					{
						continue;
					}

					if (!g1.IsZero() && f.Gradient.Contains(variable1))
					{
						expressions.AddOnce(g1);

						gradient.AddOnce(variable1);
					}

					for (int j = i; j < variables.Count; j++)
					{
						DualVariable variable2 = variables[j];

						if (!DualVariable.CheckSecondOrder(variable1, variable2))
						{
							continue;
						}

						HashedTuple<DualVariable, DualVariable> variable12 = HashedTuple.Create<DualVariable, DualVariable>(variable1, variable2);

						if (!g1.IsZero() && f.Hessian.Contains(variable12))
						{
							expressions.AddOnce(g1);

							hessian.AddOnce(variable12);
							hessian1.AddOnce(variable12);
						}

						if (!g11.IsZero() && f.Gradient.Contains(variable1) && f.Gradient.Contains(variable2))
						{
							expressions.AddOnce(g11);

							hessian.AddOnce(variable12);
							hessian11.AddOnce(variable12);
						}
					}
				}

				if (gradient == f.Gradient)
				{
					gradient = f.Gradient;
				}

				if (hessian == f.Hessian)
				{
					hessian = f.Hessian;
				}
			}

			// Add local variables to the method body.
			foreach (Expression expression in expressions)
			{
				expression.Prepare(body);
			}

			FunctionValueStorage functionValueStorage = GetFunctionValueStorage();
			Storage gradientStorage = GetStorage(gradient.Count);
			Storage hessianStorage = GetStorage(hessian.Count);

			body.AddFunctionValue(g);

			body.AddGradientLoop(functionValueStorage, gradient, gradientStorage,
				GradientLoopTerm.Create(f, g1, gradient));

			body.AddHessianLoop(functionValueStorage, hessian, hessianStorage,
				HessianLoopTerm.Create(f, g1, hessian1),
				HessianLoopTerm.CreateFromGradient(f, g11, hessian11));

			return AddCompiledFunction(function, body, functionValueStorage, gradient, gradientStorage, hessian, hessianStorage);
		}

		public GeneratedFunction AddBinaryFunction(DualFunction function, GeneratedFunction f1, GeneratedFunction f2, Expression g, Expression g1, Expression g2, Expression g11, Expression g12, Expression g22)
		{
			f1.Expression.Link("f1");
			f2.Expression.Link("f2");

			FunctionBody body = new FunctionBody(this, "g", new FunctionParameter("f1", f1), new FunctionParameter("f2", f2));

			// Replace by references to local variables.
			g1 = Expression.MarkLocal(g1);
			g2 = Expression.MarkLocal(g2);
			g11 = Expression.MarkLocal(g11);
			g12 = Expression.MarkLocal(g12);
			g22 = Expression.MarkLocal(g22);

			// List of expressions needed to compute the function value and derivatives.
			HashedList<Expression> expressions = new HashedList<Expression>();
			expressions.Add(g);

			// Build output and intermediate gradients.
			SortedGradient gradient = new SortedGradient();
			SortedGradient gradient1 = new SortedGradient();
			SortedGradient gradient2 = new SortedGradient();

			// And output and intermediate Hessians.
			SortedHessian hessian = new SortedHessian();
			SortedHessian hessian1 = new SortedHessian();
			SortedHessian hessian2 = new SortedHessian();
			SortedHessian hessian11 = new SortedHessian();
			SortedHessian hessian22 = new SortedHessian();
			SortedHessian hessian12 = new SortedHessian();
			SortedHessian hessian21 = new SortedHessian();

			if (f1.Gradient == f2.Gradient
				&& f1.Hessian.Count == f1.Gradient.HessianCount
				&& f2.Hessian.Count == f2.Gradient.HessianCount)
			{
				if (!g1.IsZero())
				{
					expressions.AddOnce(g1);

					gradient = f1.Gradient;
					gradient1 = f1.Gradient;
					hessian = f1.Hessian;
					hessian1 = f1.Hessian;
				}

				if (!g2.IsZero())
				{
					expressions.AddOnce(g2);

					gradient = f2.Gradient;
					gradient2 = f2.Gradient;
					hessian = f2.Hessian;
					hessian2 = f2.Hessian;
				}

				if (!g11.IsZero())
				{
					expressions.AddOnce(g11);

					hessian = f1.Hessian;
					hessian11 = f1.Hessian;
				}

				if (!g22.IsZero())
				{
					expressions.AddOnce(g22);

					hessian = f2.Hessian;
					hessian22 = f2.Hessian;
				}

				if (!g12.IsZero())
				{
					expressions.AddOnce(g12);

					hessian = f1.Hessian; // = f2.Hessian
					hessian12 = hessian21 = f1.Hessian; // = f2.Hessian
				}
			}
			else
			{
				IList<DualVariable> variables = SortedGradient.Merge(f1.Gradient, f2.Gradient);

				for (int i = 0; i < variables.Count; i++)
				{
					DualVariable variable1 = variables[i];

					if (!DualVariable.CheckFirstOrder(variable1))
					{
						continue;
					}

					if (!g1.IsZero() && f1.Gradient.Contains(variable1))
					{
						expressions.AddOnce(g1);

						gradient.AddOnce(variable1);
						gradient1.AddOnce(variable1);
					}

					if (!g2.IsZero() && f2.Gradient.Contains(variable1))
					{
						expressions.AddOnce(g2);

						gradient.AddOnce(variable1);
						gradient2.AddOnce(variable1);
					}

					for (int j = i; j < variables.Count; j++)
					{
						DualVariable variable2 = variables[j];

						if (!DualVariable.CheckSecondOrder(variable1, variable2))
						{
							continue;
						}

						HashedTuple<DualVariable, DualVariable> variable12 = HashedTuple.Create<DualVariable, DualVariable>(variable1, variable2);

						if (!g1.IsZero() && f1.Hessian.Contains(variable12))
						{
							expressions.AddOnce(g1);

							hessian.AddOnce(variable12);
							hessian1.AddOnce(variable12);
						}

						if (!g2.IsZero() && f2.Hessian.Contains(variable12))
						{
							expressions.AddOnce(g2);

							hessian.AddOnce(variable12);
							hessian2.AddOnce(variable12);
						}

						if (!g11.IsZero() && f1.Gradient.Contains(variable1) && f1.Gradient.Contains(variable2))
						{
							expressions.AddOnce(g11);

							hessian.AddOnce(variable12);
							hessian11.AddOnce(variable12);
						}

						if (!g22.IsZero() && f2.Gradient.Contains(variable1) && f2.Gradient.Contains(variable2))
						{
							expressions.AddOnce(g22);

							hessian.AddOnce(variable12);
							hessian22.AddOnce(variable12);
						}

						if (!g12.IsZero() && f1.Gradient.Contains(variable1) && f2.Gradient.Contains(variable2))
						{
							expressions.AddOnce(g12);

							hessian.AddOnce(variable12);
							hessian12.AddOnce(variable12);
						}

						if (!g12.IsZero() && f2.Gradient.Contains(variable1) && f1.Gradient.Contains(variable2))
						{
							expressions.AddOnce(g12);

							hessian.AddOnce(variable12);
							hessian21.AddOnce(variable12);
						}
					}
				}

				if (gradient == f1.Gradient)
				{
					gradient = f1.Gradient;
				}
				else if (gradient == f2.Gradient)
				{
					gradient = f2.Gradient;
				}

				if (hessian == f1.Hessian)
				{
					hessian = f1.Hessian;
				}
				else if (hessian == f2.Hessian)
				{
					hessian = f2.Hessian;
				}
			}

			// Add local variables to the method body.
			foreach (Expression expression in expressions)
			{
				expression.Prepare(body);
			}

			FunctionValueStorage functionValueStorage = GetFunctionValueStorage();
			Storage gradientStorage = GetStorage(gradient.Count);
			Storage hessianStorage = GetStorage(hessian.Count);

			body.AddFunctionValue(g);

			body.AddGradientLoop(functionValueStorage, gradient, gradientStorage,
				GradientLoopTerm.Create(f1, g1, gradient1),
				GradientLoopTerm.Create(f2, g2, gradient2));

			body.AddHessianLoop(functionValueStorage, hessian, hessianStorage,
				HessianLoopTerm.Create(f1, g1, hessian1),
				HessianLoopTerm.Create(f2, g2, hessian2),
				HessianLoopTerm.CreateFromGradient(f1, g11, hessian11),
				HessianLoopTerm.CreateFromGradient(f2, g22, hessian22),
				HessianLoopTerm.CreateFromGradient(f1, f2, g12, hessian12),
				HessianLoopTerm.CreateFromGradient(f2, f1, g12, hessian21));

			return AddCompiledFunction(function, body, functionValueStorage, gradient, gradientStorage, hessian, hessianStorage);
		}

		public GeneratedFunction AddCompiledFunction(DualFunction function, FunctionBody body, FunctionValueStorage functionValueStorage, SortedGradient gradient, Storage gradientStorage, SortedHessian hessian, Storage hessianStorage)
		{
			GeneratedFunction compiledFunction = new GeneratedFunction(function, functionValueStorage, gradient, gradientStorage, hessian, hessianStorage);

			//FIXME removed: compiledMapping.Add(function, compiledFunction); // FIXME test

			string method, parameters;
			body.GenerateCodeNEW(functionValueStorage, out method, out parameters);
			AddCompiledFunction(function, method, parameters);

			return compiledFunction;
		}

		private void AddCompiledFunction(DualFunction function, string method, string parameters)
		{
			string identifier;
			if (!methods.TryGetValue(method, out identifier))
			{
				identifier = "Evaluate" + nextMethod++;
				methods.Add(method, identifier);
			}

			statements.Add(identifier + "(" + parameters + ");");

			if (statements.Count % 100 == 0)
			{
				Console.WriteLine(statements.Count + "\t" + (100.0 * statements.Count / 21459).ToString("f1") + "%");
			}
		}

		private string GenerateCodeNEW(DualFunction function, out string identifier)
		{
			GetCompiledFunction(function);

			StringBuilder sb = new StringBuilder();

			sb.AppendLine("using System;");
			sb.AppendLine("using System.Collections.Generic;");
			sb.AppendLine("public class GeneratedFunction");
			sb.AppendLine("{");
			sb.AppendLine("    private Dictionary<int, double> x;");
			sb.AppendLine("    private Dictionary<string, double> y;");
			sb.AppendLine("    private double[] ax = new double[" + nextStorage + "];");
			sb.AppendLine("    private List<int[]> z;");
			sb.AppendLine();
			sb.AppendLine("    private void Prepare()"); // static constructor
			sb.AppendLine("    {");
			sb.AppendLine("        z = new List<int[]>();");
			for (int i = 0; i < offsetsSorted.Count; i++)
			{
				sb.AppendLine("        z.Add(" + offsetsSorted[i].GenerateCode() + ");");
			}
			//for (int i = 0; i < indexLists.Count; i++)
			//{
			//    sb.AppendLine("        z.Add(" + indexLists[i].GenerateCode() + ");");
			//}
			//sb.AppendLine("        z.Add(new int[] { 1, 2, 3 });");
			sb.AppendLine("    }");
			sb.AppendLine();
			sb.AppendLine("    public unsafe void Value(Dictionary<int, double> x, Dictionary<string, double> y)");
			sb.AppendLine("    {");
			sb.AppendLine("        this.x = x;");
			sb.AppendLine("        this.y = y;");
			sb.AppendLine();
			sb.AppendLine("        Prepare();");
			sb.AppendLine("        fixed (double* a = ax)");
			sb.AppendLine("        {");
			foreach (string statement in statements)
			{
				string s = statement.Replace(")", ", a)");

				sb.AppendLine("        " + s);
			}
			//sb.AppendLine("        double[] a = ax;");
			identifier = compiledMapping[function].Identifier;
			sb.AppendLine("        y[\"" + identifier + "\"] = a[" + compiledMapping[function].FunctionValueStorage.Index + "];");
			for (int i = 0; i < compiledMapping[function].Gradient.Count; i++)
			{
				DualVariable variable1 = compiledMapping[function].Gradient[i];
				sb.AppendLine("        y[\"" + identifier + "_" + variable1.FunctionIndex + "\"] = a[" + (compiledMapping[function].GradientStorage.Index + i) + "];");
			}
			for (int i = 0; i < compiledMapping[function].Hessian.Count; i++)
			{
				DualVariable variable1 = compiledMapping[function].Hessian[i].Item1;
				DualVariable variable2 = compiledMapping[function].Hessian[i].Item2;
				sb.AppendLine("        y[\"" + identifier + "_" + variable1.FunctionIndex + "_" + variable2.FunctionIndex + "\"] = a[" + (compiledMapping[function].HessianStorage.Index + i) + "];");
			}
			//sb.AppendLine("        foreach (string identifier in y_.Keys) { y[identifier] = y_[identifier]; }");
			sb.AppendLine("        }");
			sb.AppendLine("    }");
			sb.AppendLine();

			foreach (KeyValuePair<string, string> method in methods)
			{
				string methodBody = method.Key;
				string methodName = method.Value;

				//int i1 = methodBody.IndexOf("{");
				//methodBody = methodBody.Insert(i1 + 1, " fixed (double* a = ax) {");

				//int i2 = methodBody.LastIndexOf("}");
				//methodBody = methodBody.Insert(i2 + 1, "}");

				int i3 = methodBody.IndexOf(")");
				methodBody = methodBody.Insert(i3, ", double* a");

				sb.Append("    private unsafe void " + methodName + methodBody.Replace("\r\n", "\r\n    ") + "\r\n");
			}

			sb.Remove(sb.Length - 6, 6);
			sb.AppendLine("}");

			return sb.ToString();
		}

		public DualEvaluator Compile()
		{
			// Prepare the compiler.
			CompilerParameters parameters = new CompilerParameters();
			parameters.GenerateInMemory = true;
			parameters.TreatWarningsAsErrors = false;
			parameters.GenerateExecutable = false;
			parameters.CompilerOptions = "/optimize";
			parameters.IncludeDebugInformation = false;
			parameters.ReferencedAssemblies.Add("System.dll");
			parameters.CompilerOptions = "/unsafe";

			//Console.Write("Generating code: ");
			//Stopwatch sw = Stopwatch.StartNew();
			//string code = GenerateCode(f);
			//Console.WriteLine(sw.Elapsed);

			//System.IO.File.WriteAllText(@"C:\Temp\Code.txt", generatedCode);

			// Run the compiler.
			//Console.Write("Compiling: ");
			//sw = Stopwatch.StartNew();
			CompilerResults results = new CSharpCodeProvider().CompileAssemblyFromSource(parameters, generatedCode);
			//Console.WriteLine(sw.Elapsed);

			if (results.Errors.HasErrors)
			{
				// Something's wrong with the code. Just show the first error.
				string s = "Compile error: ";
				foreach (CompilerError error in results.Errors)
				{
					if (!error.IsWarning)
					{
						s += error.ToString();
						break;
					}
				}
				throw new Exception(s);
			}

			Type type = results.CompiledAssembly.GetModules()[0].GetType("GeneratedFunction");
			return new DualEvaluator(Activator.CreateInstance(type), type.GetMethod("Value"), identifier);
		}

		public string GeneratedCode
		{
			get
			{
				return generatedCode;
			}
		}

		//public string Identifier
		//{
		//    get
		//    {
		//        return identifier;
		//    }
		//}
	}
}
